#!/usr/bin/env python

import nltk
import random
from nltk.corpus import movie_reviews
from nltk import FreqDist
from nltk.corpus import stopwords

docs = [(list(movie_reviews.words(fileid)),category)
	for category in movie_reviews.categories()
	for fileid in movie_reviews.fileids(category)
]

random.shuffle(docs)
en_stopwords = stopwords.words('english')
all_words = FreqDist(w.lower() for w in movie_reviews.words() if w.isalpha() and w.lower() not in en_stopwords)
# word_sample = all_words.keys()[:2000]
word_sample = [w for (w,count) in all_words.most_common(2000)]

def doc_features(doc):
	doc_words = set(doc)
	features = {}
	for word in word_sample:
		# features['contains({})'.format(word)] = (word in doc_words)
		features[word] = (word in doc_words)
	return features

# print doc_features(movie_reviews.words('pos/cv957_8737.txt'))
# {'contains(waste)': False, 'contains(lot)': False, ...}


featuresets = [(doc_features(d),c) for (d,c) in docs]
train_set, test_set = featuresets[100:],featuresets[:100]
classifier = nltk.NaiveBayesClassifier.train(train_set);

print(nltk.classify.accuracy(classifier, test_set))
classifier.show_most_informative_features(30)